#include <sys/types.h>
#include <sys/wait.h>
#include <sys/user.h>
#include <sys/reg.h>
#include <sys/ptrace.h>
#include <sys/stat.h>
#include <sys/signal.h>
#include <sys/time.h>
#include <sys/personality.h>
#include <syscall.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdio.h>
#include <string.h>
#include "ptrace_utils.h"

#define NAME_SIZE 300
#define STACK_TOP 0xBFFFFFFF // TODO: в обще говоря, сомнительный момент

int g_stkfd;
int g_ctxfd;
pid_t g_pid;
bool g_restore;
timeval g_interval;

bool is_correct_files()
{
    return g_stkfd >= 0 && g_ctxfd >= 0;
}

int restore_state()
{
    if (!g_restore)
        return 0;
    errno = 0;


    // Восстанавливаем контекст
    struct user_regs_struct regs;
    if (read(g_ctxfd, &regs, sizeof(user_regs_struct)) <= 0) {
        perror("read context file");
        return errno;
    }

    if (trace_set_regs(g_pid, &regs) != 0)
        return errno;

    // Восстанавливаем стек
    long stack_size = STACK_TOP - regs.esp;
    printf("stack size: %08lx\n", stack_size);

    void* buf = malloc(stack_size);
    if (read(g_stkfd, buf, stack_size) <= 0) {
        perror("read stack file");
        goto END;
    }

    if (trace_write_data(g_pid, regs.esp, buf, stack_size) != 0)
        goto END;

    // TODO: Тут еще что-нибудь восстановливаем

END:
    free(buf);
    return errno;
}

int save_state()
{
    errno = 0;

    // Сохраням контекст
    struct user_regs_struct regs;
    if (trace_get_regs(g_pid, &regs) != 0)
        return errno;
    if (lseek(g_ctxfd, 0, SEEK_SET) < 0) {
        perror("lseek context file");
        return errno;
    }

    if (write(g_ctxfd, &regs, sizeof(user_regs_struct)) <= 0) {
        perror("write context file");
        return errno;
    }

    // Сохраняем стек
    long stack_size = STACK_TOP - regs.esp;
    printf("stack size: %08lx\n", stack_size);

    void* buf = malloc(stack_size);
    if (trace_read_data(g_pid, regs.esp, buf, stack_size) != 0)
        goto END;

    if (lseek(g_stkfd, 0, SEEK_SET) < 0) {
        perror("lseek stack file");
        goto END;
    }

    if (write(g_stkfd, buf, stack_size) <= 0) {
        perror("write stack file");
        goto END;
    }

    // TODO: Тут еще что-нибудь сохраняем

END:
    free(buf);
    return errno;
}

/* Выполняется родительским процессом.
 * Реализует цикл сохранения состояния дочернего процесса */
int do_parent()
{
    int status;
    siginfo_t siginfo;
    bool running = true, restored = false;
    errno = 0;

    // Ждем пока всякие exec отработают
    if (waitpid(g_pid, &status, 0) < 0) {
        perror("waitpid");
        return errno;
    }

    // Цикл сохранения состояния
    while (running) {
        trace_cont(g_pid, &status);

        // WIFEXITED возвращает ненулевое значение, если порожденный процесс был
        // завершен посредством вызова exit;
        // WIFSIGNALED возвращает 1, если выполнение процесса было завершено сигналом
        if (WIFEXITED (status) || WIFSIGNALED (status))
            break;

        if (trace_get_siginfo(g_pid, &siginfo) != 0)
            break;

        switch (siginfo.si_signo) {
            // Сработал таймер в режиме задачи
            case SIGVTALRM:
                if (!restored && g_restore) {
                    /* TODO: Нужно запускать восстановление сразу после exec */
                    restored = true;
                    running = restore_state() == 0;
                } else {
                    running = save_state() == 0;
                }
                break;

            // TODO: тут могут быть еще какие нибудь сигналы
        }
    }

    return errno;
}

/* Выполняется дочерним процессом.
 * Устанавливает флаг отладки; устанавливает таймер,
 * по которому будет реализовываться сохранение состояния;
 * загружает образ исполняемого файла в свое адресное пространство */
int do_child(char* imgname)
{
    struct itimerval alarm_timer;
    errno = 0;

    alarm_timer.it_interval.tv_sec  = g_interval.tv_sec;
    alarm_timer.it_interval.tv_usec = g_interval.tv_usec;
    alarm_timer.it_value.tv_sec     = 0;
    alarm_timer.it_value.tv_usec    = 10;

    // ITIMER_VIRTUAL - таймер уменьшается, только когда процесс выполняется в режиме задачи.
    // По истечению временного интервала отправляется сигнал SIGVTALRM
    if (setitimer(ITIMER_VIRTUAL, &alarm_timer, NULL)) {
        perror("setitimer");
        return errno;
    }
    // Отключаем рандомизацию адресов (таких как верхушка стека, например)
    personality(ADDR_NO_RANDOMIZE);

    ptrace(PTRACE_TRACEME, 0, NULL, NULL);
    execlp(imgname, imgname, NULL);

    // Сюда не должны попасть при нормальном выполнении execlp
    perror("exec");
    return errno;
}

/* st_state - программа сохранения и восстановления состояния выполняющегося процесса.
 * Входные параметры:
 *  argv[1] - имя образа исполняемого файла
    TODO: argv[2] - длинна временного интервала между сохранениями (в секундах) */
int main(int argc, char *argv[])
{
    char imgname[NAME_SIZE];
    char stkname[NAME_SIZE];
    char ctxname[NAME_SIZE];

    g_stkfd = g_ctxfd = g_pid = -1;
    g_restore = false;
    g_interval.tv_sec = 5;
    g_interval.tv_usec = 0;

    printf("SAVE-RESTORE PROCESS STATE\n");
    if (argc == 1) {
        printf("image name: ");
        scanf("%300s", imgname);
    } else {
        memcpy(imgname, argv[1], strlen(argv[1]) * sizeof(char));
    }

    // Проверка корректности образа файла
    struct stat s;
    stat(imgname, &s);
    bool is_exec = (s.st_mode & S_IXUSR) != 0 || (s.st_mode & S_IXOTH) != 0 || (s.st_mode & S_IXGRP) != 0;
    if (!S_ISREG(s.st_mode) || !is_exec) {
        printf("invalid image file\n");
        return 0;
    }

    // Получаем дескрипторы файлов контекста и стека
    sprintf(stkname, "%s_stk.bin", imgname);
    sprintf(ctxname, "%s_ctx.bin", imgname);

    mode_t mode = S_IRUSR | S_IWUSR;
    g_stkfd = open(stkname, O_RDWR, mode);
    g_ctxfd = open(ctxname, O_RDWR, mode);
    if (!is_correct_files()) {
        g_stkfd = creat(stkname, mode);
        g_ctxfd = creat(ctxname, mode);
        if (!is_correct_files()) {
            perror("can't create state files");
            return 0;
        }
    } else {
        printf("state files was opened...\n");
        g_restore = true;
    }

    // Запускаем новый процесс
    if ((g_pid = fork()) < 0) {
        perror("can't create new process");
        goto END;
    }

    if (g_pid == 0) {
        do_child(imgname);
    } else {
        do_parent();
    }

END:
    close(g_ctxfd);
    close(g_stkfd);
    return 0;
}
